// Copyright (c) 2024 Go-Frame-Lite
// Go-Frame-Lite is licensed under Mulan PSL v2.
// You can use this software according to the terms and conditions of the Mulan
// PSL v2.
// You may obtain a copy of Mulan PSL v2 at:
//         http://license.coscl.org.cn/MulanPSL2
// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY
// KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO
// NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
// See the Mulan PSL v2 for more details.

// Package trace provides distributed tracing functionality using OpenTelemetry.
package trace

import (
	"context"
	"fmt"
	"log"

	"go-frame-lite/pkg/configtypes"

	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/codes"
	"go.opentelemetry.io/otel/exporters/stdout/stdouttrace"
	"go.opentelemetry.io/otel/propagation"
	"go.opentelemetry.io/otel/sdk/resource"
	sdktrace "go.opentelemetry.io/otel/sdk/trace"
	semconv "go.opentelemetry.io/otel/semconv/v1.4.0"
	"go.opentelemetry.io/otel/trace"
	"go.opentelemetry.io/otel/trace/noop"
)

// TracerProvider 全局追踪器提供者
var TracerProvider trace.TracerProvider

// Tracer 全局追踪器
var Tracer trace.Tracer

// DefaultTraceHeader 默认的TraceID请求头
const DefaultTraceHeader = "traceparent"

// DefaultSpanHeader 默认的SpanID请求头
const DefaultSpanHeader = "tracestate"

// InitTracer 初始化OpenTelemetry追踪器
func InitTracer(config configtypes.TraceConfig) error {
	if !config.Enabled {
		// 如果追踪未启用，使用无操作追踪器
		TracerProvider = noop.NewTracerProvider()
		Tracer = TracerProvider.Tracer(config.ServiceName)
		return nil
	}

	// 创建标准输出导出器（生产环境可替换为Jaeger、Zipkin等）
	exporter, err := stdouttrace.New(stdouttrace.WithPrettyPrint())
	if err != nil {
		return fmt.Errorf("创建追踪导出器失败: %w", err)
	}

	// 创建资源
	res, err := resource.New(context.Background(),
		resource.WithAttributes(
			semconv.ServiceNameKey.String(config.ServiceName),
			attribute.String("environment", "production"),
		),
	)
	if err != nil {
		return fmt.Errorf("创建资源失败: %w", err)
	}

	// 创建追踪器提供者
	TracerProvider = sdktrace.NewTracerProvider(
		sdktrace.WithBatcher(exporter),
		sdktrace.WithResource(res),
		sdktrace.WithSampler(sdktrace.AlwaysSample()),
	)

	// 设置全局追踪器提供者
	otel.SetTracerProvider(TracerProvider)

	// 设置传播器（支持W3C Trace Context标准）
	propagator := propagation.NewCompositeTextMapPropagator(
		propagation.TraceContext{},
		propagation.Baggage{},
	)
	otel.SetTextMapPropagator(propagator)

	// 创建追踪器
	Tracer = TracerProvider.Tracer(config.ServiceName)

	log.Printf("OpenTelemetry追踪器已初始化，服务名: %s", config.ServiceName)
	return nil
}

// ShutdownTracer 关闭追踪器
func ShutdownTracer() error {
	if tp, ok := TracerProvider.(*sdktrace.TracerProvider); ok {
		return tp.Shutdown(context.Background())
	}
	return nil
}

// StartSpan 开始一个新的Span
func StartSpan(ctx context.Context, name string, opts ...trace.SpanStartOption) (context.Context, trace.Span) {
	return Tracer.Start(ctx, name, opts...)
}

// GetTraceID 从context中获取TraceID
func GetTraceID(ctx context.Context) string {
	span := trace.SpanFromContext(ctx)
	if span.SpanContext().IsValid() {
		return span.SpanContext().TraceID().String()
	}
	return ""
}

// GetSpanID 从context中获取SpanID
func GetSpanID(ctx context.Context) string {
	span := trace.SpanFromContext(ctx)
	if span.SpanContext().IsValid() {
		return span.SpanContext().SpanID().String()
	}
	return ""
}

// AddSpanAttribute 为当前Span添加属性
func AddSpanAttribute(ctx context.Context, key string, value interface{}) {
	span := trace.SpanFromContext(ctx)
	if span.SpanContext().IsValid() {
		switch v := value.(type) {
		case string:
			span.SetAttributes(attribute.String(key, v))
		case int:
			span.SetAttributes(attribute.Int(key, v))
		case int64:
			span.SetAttributes(attribute.Int64(key, v))
		case bool:
			span.SetAttributes(attribute.Bool(key, v))
		case float64:
			span.SetAttributes(attribute.Float64(key, v))
		default:
			span.SetAttributes(attribute.String(key, fmt.Sprintf("%v", v)))
		}
	}
}

// RecordSpanError 记录Span错误
func RecordSpanError(ctx context.Context, err error, description ...string) {
	span := trace.SpanFromContext(ctx)
	if span.SpanContext().IsValid() {
		if len(description) > 0 {
			span.RecordError(err, trace.WithAttributes(
				attribute.String("error.description", description[0]),
			))
		} else {
			span.RecordError(err)
		}
		span.SetStatus(codes.Error, err.Error())
	}
}

// ExtractTraceContext 从HTTP请求头中提取追踪上下文
func ExtractTraceContext(ctx context.Context, headers map[string][]string) context.Context {
	propagator := otel.GetTextMapPropagator()
	if propagator != nil {
		return propagator.Extract(ctx, propagation.HeaderCarrier(headers))
	}
	return ctx
}

// InjectTraceContext 将追踪上下文注入到HTTP请求头中
func InjectTraceContext(ctx context.Context, headers map[string][]string) {
	propagator := otel.GetTextMapPropagator()
	if propagator != nil {
		propagator.Inject(ctx, propagation.HeaderCarrier(headers))
	}
}
